#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Author: yzy1

import os
import sys
import shutil
import argparse
import colorama

PROGRAM_NAME = "aspirn-cli"
RUN_COMMAND = {
    "c": "gcc '$FILE$' -o '$FILE_WITHOUT_EXT$' && './$FILE_WITHOUT_EXT$'",
    "cpp":
    "g++ -std=c++11 '$FILE$' -o '$FILE_WITHOUT_EXT$' && './$FILE_WITHOUT_EXT$'",
    "cc":
    "g++ -std=c++11 '$FILE$' -o '$FILE_WITHOUT_EXT$' && './$FILE_WITHOUT_EXT$'",
    "java": "javac '$FILE$' && java '$FILE_WITHOUT_EXT$'",
    "py": "python3 '$FILE$'",
}
DEFAULT_RUN_COMMAND = "'./$FILE$'"
args = None


def ParseArgs():
    global args
    parser = argparse.ArgumentParser(PROGRAM_NAME)
    parser.add_argument(
        "-s",
        "--std",
        action="store",
        dest="stdFile",
        help="standard code or executable file",
    )
    parser.add_argument(
        "-i",
        "--input-file",
        action="store",
        dest="inputFiles",
        nargs="+",
        help="input file(s)",
    )
    parser.add_argument(
        "-g",
        "--generator",
        action="append",
        dest="generators",
        nargs="+",
        help=
        "generator code or executable file, file count of each generator, args",
    )
    parser.add_argument(
        "-o",
        "--output-dir",
        action="store",
        dest="outputDir",
        help="output folder",
        default="data",
    )
    parser.add_argument(
        "-f",
        "--force",
        action="store_true",
        dest="force",
        help="delete all files in the output folder without asking",
        default=False,
    )
    parser.add_argument(
        "-d",
        "--show-details",
        action="store_true",
        dest="showDetails",
        help="show details",
        default=False,
    )
    args = parser.parse_args()
    return


def PrintError(msg: str):
    print(colorama.Fore.RED + "Error:", colorama.Fore.RESET + msg)
    return


def PrintDetail(title: str, detail: str = ""):
    global args
    if args.showDetails:
        print(colorama.Fore.YELLOW + title, colorama.Fore.RESET + detail)
    return


def PrintTitle(title: str):
    global args
    if args.showDetails:
        print(colorama.Fore.CYAN + title)
    return


# 获取编译运行命令
def GetRunCommand(path: str) -> str:
    ext = os.path.splitext(path)[-1][1:]  # 扩展名
    return (RUN_COMMAND[ext] if ext in RUN_COMMAND.keys() else
            DEFAULT_RUN_COMMAND).replace("$FILE$", path).replace(
                "$FILE_WITHOUT_EXT$", "".join(os.path.splitext(path)[:-1]))


def RunCommandOrExit(command: str):
    returnValue = os.system(command)
    if returnValue != 0:
        exit(returnValue)
    return


def GetGenerateCommands(fullAuto: bool):
    global args
    generateCommands = []
    willGenerateFiles = []
    for f in args.inputFiles if args.inputFiles else ():
        willGenerateFile = "%d.in" % len(generateCommands)
        willGeneratePath = "%s/%s" % (args.outputDir, willGenerateFile)
        command = "cat '%s' > '%s' || type '%s' > '%s'" % (f, willGeneratePath,
                                                           f, willGeneratePath)
        generateCommands.append(command)
        willGenerateFiles.append(willGeneratePath)
    for g in args.generators if args.generators else ():
        if len(g) > 3:
            PrintError(
                "Parameter '--generator' requires 1 to 3 values, find %d" %
                len(g))
            print(g)
            exit(1)
        gPath = g[0]
        gCount = int(g[1]) if len(g) >= 2 else 1
        gArgs = g[2] if len(g) >= 3 else ""
        for _ in range(gCount):
            willGenerateFile = "%d.in" % len(generateCommands)
            willGeneratePath = "%s/%s" % (args.outputDir, willGenerateFile)
            command = "%s %s > '%s'" % (GetRunCommand(
                os.path.relpath(gPath)), gArgs, willGeneratePath)
            generateCommands.append(command)
            willGenerateFiles.append(willGeneratePath)
    return generateCommands, willGenerateFiles


def GenerateInput(generateCommands, inputFiles):
    for i, g in enumerate(generateCommands):
        PrintDetail("- %s:" % inputFiles[i], g)
        RunCommandOrExit(g)
    return


def GenerateOutput(stdCommand, inputFiles):
    global args
    if not inputFiles:
        outputPath = "%s/0.out" % (args.outputDir)
        command = "%s > %s" % (stdCommand, outputPath)
        PrintDetail("- %s:" % outputPath, command)
        RunCommandOrExit(command)
    else:
        for i, f in enumerate(inputFiles):
            outputPath = "%s/%d.out" % (args.outputDir, i)
            command = "%s < %s > %s" % (stdCommand, f, outputPath)
            PrintDetail("- %s:" % outputPath, command)
            RunCommandOrExit(command)
    return


def Main():
    global args
    colorama.init(autoreset=True)
    ParseArgs()

    if not os.path.exists(args.outputDir):
        os.makedirs(args.outputDir)
    elif args.force:
        shutil.rmtree(args.outputDir)
        os.mkdir(args.outputDir)
    elif os.listdir(args.outputDir):
        PrintError("Output folder is not empty")
        exit(1)

    fullAuto = (not args.stdFile) and (not args.inputFiles) and (
        not args.generators)

    if fullAuto:
        files = os.listdir(".")
        stdFiles = [
            f for i, f in enumerate(files)
            if f.endswith(tuple(RUN_COMMAND.keys())) != -1 and f.startswith((
                "main", "std"))
        ]
        if stdFiles:
            args.stdFile = stdFiles[0]
        else:
            PrintError("Can not find std file")
        args.inputFiles = [f for i, f in enumerate(files) if f.endswith(".in")]
        args.generators = [
            [f] for i, f in enumerate(files)
            if f.endswith(tuple(RUN_COMMAND.keys())) and f.startswith((
                "data", "make", "mkdata", "gen"))
        ]
    stdFile = os.path.relpath(args.stdFile)
    stdCommand = GetRunCommand(stdFile)
    PrintDetail("Std file:", stdFile)
    PrintDetail("Std command:", stdCommand)
    generateInputCommands, inputFiles = GetGenerateCommands(fullAuto)

    PrintTitle("Generating input files...")
    GenerateInput(generateInputCommands, inputFiles)
    PrintTitle("Generating output files...")
    GenerateOutput(stdCommand, inputFiles)
    return


if __name__ == "__main__":
    Main()
